import cv2
from PIL import Image
import numpy as np
import os



def extract_frames(video_file, Png_Name,output_folder):
    video = cv2.VideoCapture(video_file, cv2.CAP_FFMPEG)
    if not video.isOpened():
        print("无法打开视频文件")
        return
    frame_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    output_file_format = output_folder + "/" + Png_Name + "_{:03d}.png"
    frame_count = 0
    while True:
        ret, frame = video.read()
        if not ret:
            break
        output_file = output_file_format.format(frame_count)
        # 调整图像尺寸为视频帧的尺寸
        frame = cv2.resize(frame, (frame_width, frame_height))
        #frame = cv2.resize ( frame , (1920 , 1080) )
        cv2.imwrite(output_file, frame)
        frame_count += 1
    video.release()
    print("帧提取完成")

def RemoveBlackBackground ( input_path: str , file_name: str , output_path: str ) :
    target_im = Image.open ( os.path.join ( input_path , file_name ) )
    target_im = target_im.convert ( 'RGBA' )
    target_img_array = np.array ( target_im )
    # 透明度
    target_img_Alpha = np.max ( target_img_array [ : , : , :3 ] , axis = 2 )
    target_img_array [ : , : , 3 ] = target_img_Alpha
    # 保证在黑色背景上时与原图一致
    target_img_Alpha = target_img_Alpha / 255
    target_img_array [ : , : , :3 ] / target_img_Alpha.reshape ( [ target_img_array.shape [ 0 ] , -1 , 1 ] )
    res_img = Image.fromarray ( target_img_array )
    if not file_name.endswith ( ".png" ) :
        res_img.save ( os.path.join ( output_path , file_name [ :-4 ] + ".png" ) )
    else :
        res_img.save ( os.path.join ( output_path , file_name ) )


# 示例用法

input_path  = "D:/uuuuuuuu/"
output_folder = "D:/uuuuuuuu/ww"
for filename in os.listdir(input_path):
    if filename.endswith ( ".mov" ) or filename.endswith(".mp4") :
        video_input_path = os.path.join(input_path,filename).replace("\\","/")
        png_folderName= filename.split (".")[0]
        png_output_folder = os.path.join(output_folder,png_folderName).replace("\\","/")
        if not os.path.exists ( png_output_folder ) :
            os.mkdir ( png_output_folder )
        extract_frames ( video_input_path , png_folderName,png_output_folder )
        
        
        # 使用RemoveBlackBackground函数将提取的帧转换为带透明通道的PNG
        for filename in os.listdir ( png_output_folder ) :
            if filename.endswith ( ".png" ) :
                RemoveBlackBackground ( png_output_folder , filename , png_output_folder )
        print ( "PNG修改完成" )
